# -*- coding: utf-8 -*-
'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
#作者：cacho_37967865
#博客：https://blog.csdn.net/sinat_37967865
#文件：pycryptodome_model.py
#日期：2019-07-31
#备注：多种加解密方法    # pip install pycryptodome 
用pyCryptodome模块带的aes先将秘钥以及要加密的文本填充为16位   AES key must be either 16, 24, or 32 bytes long
https://pycryptodome.readthedocs.io/en/latest/src/examples.html#encrypt-data-with-aes
'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''

import base64
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP


class encryption():

    def __init__(self,count):
        self.count = count

    #  bytes不是32的倍数那就补足为count的倍数
    def add_to_count(self,value):
        while len(value) % self.count != 0:
            value += b'\x00'
        return value     # 返回bytes


    # str转换为bytes超过count位时处理
    def get_bytes(self,org_str):
        org_bytes = str.encode(org_str)
        n = int(len(org_bytes) / self.count)
        #print('bytes长度：',len(org_bytes))
        i = 0
        new_bytes = b''
        while n >= 1:
            i = i + 1
            new_byte = org_bytes[(i-1)*self.count:self.count*i-1]
            new_bytes += new_byte
            n = n - 1
        if len(org_bytes) % self.count == 0:                   # 如果是count的倍数，直接取值
            all_bytes = org_bytes
        elif len(org_bytes) % self.count != 0 and n>1:         # 如果不是count的倍数，每次截取count位相加，最后再加剩下的并补齐count位
            all_bytes = new_bytes + self.add_to_count (org_bytes[i*self.count:])
        else:
            all_bytes = self.add_to_count (org_bytes)          # 如果不是count的倍数，并且小于count位直接补齐
        #print(all_bytes)
        return all_bytes


    def AES_encrypt(self,org_str,key):
        print('--------------->加密开始<---------------')
        # 初始化加密器
        aes = AES.new(self.get_bytes(key), AES.MODE_ECB)
        #先进行aes加密
        encrypt_aes = aes.encrypt(self.get_bytes(org_str))
        # 用base64转成字符串形式
        encrypted_text = str(base64.encodebytes(encrypt_aes), encoding='utf-8')  # 执行加密并转码返回bytes
        print('加密后的数据：',encrypted_text)
        return(encrypted_text)


    def AES_decrypt(self,secret_str,key):
        print('--------------->解密开始<---------------')
        # 初始化加密器
        aes = AES.new(self.get_bytes(key), AES.MODE_ECB)
        # 优先逆向解密base64成bytes
        base64_decrypted = base64.decodebytes(secret_str.encode(encoding='utf-8'))
        # 执行解密密并转码返回str
        decrypted_text = str(aes.decrypt(base64_decrypted), encoding='utf-8').replace('\0', '')
        print('解密后的数据：',decrypted_text)


    # 生成公钥和私钥
    def generate_key(self):
        key = RSA.generate(1024)     # 256倍数并且>= 1024
        # print(key)
        # 提取私钥并存入文件
        private_key = key.exportKey()
        file_out = open("private_key.pem", "wb")
        file_out.write(private_key)

        # 提取公钥存入文件
        public_key = key.publickey().exportKey()
        file_out = open("public_key.pem", "wb")
        file_out.write(public_key)

    # 加密/过程(无法加密太长的内容，可以分段加密)
    def RSA_encrypt(self,secret_str):
        self.generate_key()
        public_key = RSA.importKey(open("public_key.pem").read())   # 读取文件
        e_cipher = PKCS1_OAEP.new(public_key)                         # 实例化加密套件
        encrypted_data = e_cipher.encrypt(secret_str.encode())        # 加密
        print(base64.encodebytes(encrypted_data))
        # 解密
        private_key = RSA.importKey(open("private_key.pem", "rb").read())
        d_cipher = PKCS1_OAEP.new(private_key)
        org_data = d_cipher.decrypt(encrypted_data)
        print(org_data.decode())


if __name__ == '__main__':
    ways = encryption(16)                           # 加密位数
    #org_str = 'http://mp.weixin.qq.com/s?__biz=MjM5NjAxOTU4MA==&amp;mid=3009217590&amp;idx=1&amp;sn=14532c49bc8cb0817544181a10e9309f&amp;chksm=90460825a7318133e7905c02e708d5222abfea930e61b4216f15b7504e39734bcd41cfb0a26d&amp;scene=27#wechat_redirect'
    org_str = 'http://mp.weixin.qq.com/s?__biz=MjM5NjAxOTU4MA==&amp;mid=3009217590&amp;'
    key = '19881103'                                 # 秘钥
    #secret_str = ways.AES_encrypt(org_str,key)      # 加密
    #ways.AES_decrypt(secret_str,key)                # 解密
    ways.RSA_encrypt(org_str)